﻿namespace FATEFAKE
{
    using System;
    using System.Diagnostics;
    using System.Runtime.InteropServices;


    internal class JobHelper : IDisposable
    {
        [DllImport("kernel32.dll")]
        public static extern bool IsProcessInJob(IntPtr Process, IntPtr Job, out bool Result);

        [DllImport("kernel32.dll", CharSet = CharSet.Unicode)]
        static extern IntPtr CreateJobObject(IntPtr a, string lpName);

        [DllImport("kernel32.dll")]
        static extern bool SetInformationJobObject(IntPtr hJob, JobObjectInfoType infoType, IntPtr lpJobObjectInfo, UInt32 cbJobObjectInfoLength);

        [DllImport("kernel32.dll", SetLastError = true)]
        static extern bool AssignProcessToJobObject(IntPtr job, IntPtr process);

        [DllImport("kernel32.dll", SetLastError = true)]
        [return: MarshalAs(UnmanagedType.Bool)]
        static extern bool CloseHandle(IntPtr hObject);

        private IntPtr handle;
        private bool disposed;

        public JobHelper(Process process)
        {
            bool isAlreadyInJob;

            // If we are part of a job, then we cannot create a child which is in a different job (at least not
            // that easily), so in that case we do nothing. This happens e. g. when launch from VisualStudio
            if (IsProcessInJob(Process.GetCurrentProcess().Handle, IntPtr.Zero, out isAlreadyInJob) && isAlreadyInJob == true)
            {
                this.handle = IntPtr.Zero;
                return;
            }

            this.handle = CreateJobObject(IntPtr.Zero, null);
            if (this.handle != IntPtr.Zero)
            {
                // set the flag "JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE"
                // -> Causes all processes associated with the job to terminate when the last handle to the job is closed.
                var info = new JOBOBJECT_BASIC_LIMIT_INFORMATION { LimitFlags = 0x2000 };
                var extendedInfo = new JOBOBJECT_EXTENDED_LIMIT_INFORMATION { BasicLimitInformation = info };
                int length = Marshal.SizeOf(typeof(JOBOBJECT_EXTENDED_LIMIT_INFORMATION));
                IntPtr extendedInfoPtr = Marshal.AllocHGlobal(length);
                Marshal.StructureToPtr(extendedInfo, extendedInfoPtr, false);
                if (!SetInformationJobObject(handle, JobObjectInfoType.ExtendedLimitInformation, extendedInfoPtr, (uint)length))
                {
                    throw new Exception(string.Format("Unable to set information.  Error: {0}", GetWin32ErrorText(Marshal.GetLastWin32Error())));
                }

                this.AddProcess(process.Handle);
            }
        }

        public void Dispose()
        {
            Dispose(true);
            GC.SuppressFinalize(this);
        }

        private void Dispose(bool disposing)
        {
            if (this.disposed)
            {
                return;
            }

            this.Close();
            this.disposed = true;
        }

        public void Close()
        {
            if (this.handle != IntPtr.Zero)
            {
                CloseHandle(this.handle);
                this.handle = IntPtr.Zero;
            }
        }

        private bool AddProcess(IntPtr processHandle)
        {
            bool b = AssignProcessToJobObject(handle, processHandle);
            if (b == false)
            {
                throw new Exception(string.Format("Unable to assign process to job.  Error: {0}", GetWin32ErrorText(Marshal.GetLastWin32Error())));
            }

            return b;
        }

        private bool AddProcess(int processId)
        {
            return AddProcess(Process.GetProcessById(processId).Handle);
        }

        static private string GetWin32ErrorText(int error)
        {
            return string.Format("{0} -> {1}", error, new System.ComponentModel.Win32Exception(error).Message);
        }

    }

    #region Helper classes
    [StructLayout(LayoutKind.Sequential)]
    struct IO_COUNTERS
    {
        public UInt64 ReadOperationCount;
        public UInt64 WriteOperationCount;
        public UInt64 OtherOperationCount;
        public UInt64 ReadTransferCount;
        public UInt64 WriteTransferCount;
        public UInt64 OtherTransferCount;
    }

    [StructLayout(LayoutKind.Sequential)]
    struct JOBOBJECT_BASIC_LIMIT_INFORMATION
    {
        public Int64 PerProcessUserTimeLimit;
        public Int64 PerJobUserTimeLimit;
        public UInt32 LimitFlags;
        public UIntPtr MinimumWorkingSetSize;
        public UIntPtr MaximumWorkingSetSize;
        public UInt32 ActiveProcessLimit;
        public UIntPtr Affinity;
        public UInt32 PriorityClass;
        public UInt32 SchedulingClass;
    }

    [StructLayout(LayoutKind.Sequential)]
    public struct SECURITY_ATTRIBUTES
    {
        public UInt32 nLength;
        public IntPtr lpSecurityDescriptor;
        public Int32 bInheritHandle;
    }

    [StructLayout(LayoutKind.Sequential)]
    struct JOBOBJECT_EXTENDED_LIMIT_INFORMATION
    {
        public JOBOBJECT_BASIC_LIMIT_INFORMATION BasicLimitInformation;
        public IO_COUNTERS IoInfo;
        public UIntPtr ProcessMemoryLimit;
        public UIntPtr JobMemoryLimit;
        public UIntPtr PeakProcessMemoryUsed;
        public UIntPtr PeakJobMemoryUsed;
    }

    public enum JobObjectInfoType
    {
        AssociateCompletionPortInformation = 7,
        BasicLimitInformation = 2,
        BasicUIRestrictions = 4,
        EndOfJobTimeInformation = 6,
        ExtendedLimitInformation = 9,
        SecurityLimitInformation = 5,
        GroupInformation = 11
    }
    #endregion
}